Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

deepmd/dpmodel/utils/lmdb_data.py (5)

248-260: ⚠️ Potential issue | 🟡 Minor

Validate batch_size to prevent silent misconfiguration.

Unknown string specs (e.g., "auto_invalid") fall through to self._auto_rule = 32 at line 256, hiding config mistakes. Non-positive integers are also accepted at line 260.

Suggested fix

         self._auto_rule: int | None = None
         if isinstance(batch_size, str):
             if batch_size == "auto":
                 self._auto_rule = 32
             elif batch_size.startswith("auto:"):
-                self._auto_rule = int(batch_size.split(":")[1])
+                try:
+                    self._auto_rule = int(batch_size.split(":", 1)[1])
+                except ValueError as exc:
+                    raise ValueError(
+                        f"Invalid batch_size spec: {batch_size!r}"
+                    ) from exc
+                if self._auto_rule <= 0:
+                    raise ValueError("auto batch_size rule must be > 0")
             else:
-                self._auto_rule = 32
+                raise ValueError(
+                    f"batch_size must be 'auto', 'auto:N', or a positive int, got {batch_size!r}"
+                )
             self.batch_size = _compute_batch_size(self._natoms, self._auto_rule)
         else:
             self.batch_size = int(batch_size)
+            if self.batch_size <= 0:
+                raise ValueError("batch_size must be > 0")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 248 - 260, The batch_size
parsing in the constructor silently treats unknown string specs and non-positive
integers as valid (setting self._auto_rule = 32 or accepting non-positive ints),
so add explicit validation: in the batch_size handling for strings
(function/class using self._auto_rule and _compute_batch_size) only accept
"auto" or "auto:<positive_int>" and raise a ValueError for any other string
spec; when batch_size is numeric (the else branch that sets self.batch_size =
int(batch_size)) validate that the parsed int is > 0 and raise ValueError if
not; ensure error messages reference the provided batch_size value and keep use
of _compute_batch_size(self._natoms, self._auto_rule) unchanged when the auto
rule is valid.

---

974-976: ⚠️ Potential issue | 🟡 Minor

Guard against truncated frame_nlocs metadata during merge.

If src_nlocs has fewer entries than nframes (malformed metadata), src_nlocs[i] raises IndexError.

Suggested fix

-                if src_nlocs is not None:
+                if src_nlocs is not None and i < len(src_nlocs):
                     frame_nlocs.append(int(src_nlocs[i]))
                 else:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 974 - 976, The code indexing
src_nlocs[i] can raise IndexError when src_nlocs is shorter than nframes; in the
block that checks "if src_nlocs is not None:" (where
frame_nlocs.append(int(src_nlocs[i])) is called), first verify i <
len(src_nlocs) (or use a safe iterator) and if the entry is missing append a
sensible default (e.g., 0) or otherwise handle the malformed metadata (log/warn
and append 0) instead of indexing past the end; update the logic around
frame_nlocs, src_nlocs and nframes to perform this bounds check before
casting/indexing.

---

785-808: ⚠️ Potential issue | 🟡 Minor

Handle empty datasets in get_test().

If the LMDB has zero frames, self._nloc_groups is empty and max(self._nloc_groups, ...) at line 798 raises ValueError: max() arg is an empty sequence.

Suggested fix

         if nloc is not None:
             if nloc not in self._nloc_groups:
                 raise ValueError(
                     f"No frames with nloc={nloc}. Available: {sorted(self._nloc_groups.keys())}"
                 )
             frame_indices = self._nloc_groups[nloc]
             natoms = nloc
+        elif not self._frames:
+            raise ValueError("LMDB dataset contains no frames")
         elif len(self._nloc_groups) == 1:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 785 - 808, The method that
selects frames (inside get_test / the function using self._nloc_groups and
_frames) fails when the LMDB is empty because max(...) is called on an empty
self._nloc_groups; add an early check for empty dataset (e.g., if not
self._nloc_groups or not self._frames) and handle it by either raising a clear
ValueError ("No frames in LMDB") or returning an appropriate empty result before
reaching the mixed-nloc branch; update the branches that reference natoms and
frame_indices (the code that computes natoms = max(...) and frame_indices = ...)
so they only run when groups exist to avoid the max() on empty sequence.

---

436-447: ⚠️ Potential issue | 🟡 Minor

Remove unused unique_nlocs variable (Ruff F841).

unique_nlocs is assigned but never used. This will cause CI failure per coding guidelines.

Suggested fix

     def print_summary(self, name: str, prob: Any) -> None:
         """Print basic dataset info."""
-        unique_nlocs = sorted(self._nloc_groups.keys())
         nloc_info = ", ".join(
             f"{nloc}({len(idxs)})" for nloc, idxs in sorted(self._nloc_groups.items())
         )

As per coding guidelines, **/*.py: Always run ruff check . and ruff format . before committing changes or CI will fail.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 436 - 447, In print_summary
(method print_summary) the local variable unique_nlocs is assigned from
self._nloc_groups but never used; remove the unused assignment (unique_nlocs =
sorted(self._nloc_groups.keys())) so the method only builds nloc_info from
self._nloc_groups and logs it, eliminating the Ruff F841 unused-variable error.

---

265-276: ⚠️ Potential issue | 🟡 Minor

Guard against out-of-range atom type IDs.

np.bincount(...)[: self._ntypes] truncates counts for type IDs >= self._ntypes, making natoms_vec inconsistent with frame contents.

Suggested fix

     def _compute_natoms_vec(self, atype: np.ndarray) -> np.ndarray:
         nloc = len(atype)
+        if atype.size and (np.min(atype) < 0 or np.max(atype) >= self._ntypes):
+            raise ValueError(
+                f"atype contains out-of-range IDs; expected [0, {self._ntypes - 1}]"
+            )
         counts = np.bincount(atype, minlength=self._ntypes)[: self._ntypes]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 265 - 276,
_compute_natoms_vec currently truncates counts for atom type IDs >=
self._ntypes; to fix, validate the atype array before counting: compute max_id =
atype.max() (handle empty atype properly) and if max_id >= self._ntypes or any
id < 0 raise a clear ValueError referencing _compute_natoms_vec and include the
offending max_id and self._ntypes in the message; otherwise use
np.bincount(atype, minlength=self._ntypes) to build counts and keep the existing
vec[0]=nloc, vec[1]=nloc, vec[2:]=counts logic.

deepmd/pt/train/training.py (1)

263-265: Consider documenting the _reader access pattern.

_data._reader accesses a private attribute of LmdbDataset. While acceptable within the same package, consider either making this a public property or adding a brief comment explaining this coupling is intentional for distributed sampler construction.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt/train/training.py` around lines 263 - 265, The code is accessing a
private attribute _data._reader of LmdbDataset when constructing the
DistributedSameNlocBatchSampler; update the codebase to avoid relying on a
private member by either exposing a public property/method on LmdbDataset (e.g.,
reader or get_reader) and using that, or add a concise inline comment next to
the DistributedSameNlocBatchSampler construction documenting that accessing
_reader is an intentional coupling for distributed sampler creation and should
not be refactored without updating the sampler; reference symbols:
_data._reader, DistributedSameNlocBatchSampler, LmdbDataset.

deepmd/dpmodel/utils/lmdb_data.py (1)

869-869: Rename unused loop variable to _req_info.

req_info is not used within the loop body (Ruff B007). Rename to _req_info to indicate it's intentionally unused.

-        for key, req_info in all_keys.items():
+        for key, _req_info in all_keys.items():

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` at line 869, The loop over all_keys uses
an unused variable name req_info which triggers a lint warning; update the for
statement in lmdb_data.py (the loop "for key, req_info in all_keys.items():") to
rename req_info to _req_info so it becomes "for key, _req_info in
all_keys.items():" to indicate the variable is intentionally unused (no other
changes needed if req_info is not referenced elsewhere).

deepmd/dpmodel/utils/lmdb_data.py (2)

892-896: Rename unused loop variable req_info to _req_info (Ruff B007).

The loop variable req_info is not used within the loop body. Rename it to _req_info to signal intentional non-use and satisfy Ruff.

🧹 Proposed fix

-        for key, req_info in all_keys.items():
+        for key, _req_info in all_keys.items():
             has_key = any(
                 key in f and isinstance(f.get(key), np.ndarray) for f in frames
             )

As per coding guidelines, **/*.py: Always run ruff check . and ruff format . before committing changes or CI will fail.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 892 - 896, The loop in
lmdb_data.py iterates with an unused variable req_info (for key, req_info in
all_keys.items()) — rename it to _req_info to indicate intentional non-use and
satisfy Ruff B007; update the for-loop variable in that block (the loop that
sets result[f"find_{key}"]) and run ruff check/ruff format before committing.

---

962-966: Consider safer directory removal approach.

shutil.rmtree(dst_path) on line 966 unconditionally removes the existing path. While this enables idempotent merges, consider adding a confirmation parameter or logging a warning when overwriting to prevent accidental data loss.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 962 - 966, The code
unconditionally deletes dst_path using shutil.rmtree which can cause accidental
data loss; modify the enclosing function (look for the routine that uses
dst_path and shutil.rmtree) to accept an explicit boolean flag like overwrite or
force_remove (default False), add logging (import logging) and call
logging.warning with the dst_path before deletion, and only call
shutil.rmtree(dst_path) when the flag is True (otherwise raise an error or
return); this makes removal explicit and auditable while keeping the idempotent
behavior when explicitly requested.

deepmd/dpmodel/utils/lmdb_data.py (5)

367-371: ⚠️ Potential issue | 🟠 Major

Fail fast on invalid atype ids before remap/counting.

At Line 371 and Line 291, atype is used without bounds validation. Invalid ids can corrupt type remapping/counts or raise runtime errors in non-obvious paths. Validate range once before remap and before natoms_vec computation.

Suggested fix

         if "atype" in frame and isinstance(frame["atype"], np.ndarray):
             frame["atype"] = frame["atype"].reshape(-1).astype(np.int64)
+            if frame["atype"].size and np.min(frame["atype"]) < 0:
+                raise ValueError("atype contains negative type ids")
             # Remap atom types from LMDB's type_map to model's type_map
             if self._type_remap is not None:
+                if frame["atype"].size and np.max(frame["atype"]) >= len(self._type_remap):
+                    raise ValueError(
+                        f"atype id out of range for LMDB type_map (max allowed {len(self._type_remap)-1})"
+                    )
                 frame["atype"] = self._type_remap[frame["atype"]].astype(np.int64)
+            if frame["atype"].size and np.max(frame["atype"]) >= self._ntypes:
+                raise ValueError(
+                    f"atype contains out-of-range ids; expected [0, {self._ntypes - 1}]"
+                )

Also applies to: 290-292

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 367 - 371, Validate that all
atom type IDs in frame["atype"] are within the valid range before any remapping
or any use in natoms_vec computation: check they are non-negative integers and
less than the length/shape of self._type_remap (when self._type_remap is not
None) or less than expected num_types; if any id is out of range raise a clear
ValueError including the offending ids and context (frame id), then proceed to
reshape/astype, apply self._type_remap on valid indices, and only after
successful remap compute natoms_vec; refer to symbols frame["atype"],
self._type_remap and natoms_vec to locate where to insert the validation and
error handling.

---

270-281: ⚠️ Potential issue | 🟠 Major

Validate batch_size strictly; don’t silently coerce invalid specs.

At Line 276, unknown string specs are silently treated as auto:32, and at Line 280 non-positive integers are accepted. This masks config mistakes and can break batching downstream.

Suggested fix

         self._auto_rule: int | None = None
         if isinstance(batch_size, str):
             if batch_size == "auto":
                 self._auto_rule = 32
             elif batch_size.startswith("auto:"):
-                self._auto_rule = int(batch_size.split(":")[1])
+                try:
+                    self._auto_rule = int(batch_size.split(":", 1)[1])
+                except ValueError as exc:
+                    raise ValueError(
+                        "batch_size must be 'auto', 'auto:N', or a positive integer"
+                    ) from exc
+                if self._auto_rule <= 0:
+                    raise ValueError("auto batch_size rule N must be > 0")
             else:
-                self._auto_rule = 32
+                raise ValueError(
+                    "batch_size must be 'auto', 'auto:N', or a positive integer"
+                )
             # Default batch_size uses first frame's nloc (for total_batch estimate)
             self.batch_size = _compute_batch_size(self._natoms, self._auto_rule)
         else:
             self.batch_size = int(batch_size)
+            if self.batch_size <= 0:
+                raise ValueError("batch_size must be > 0")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 270 - 281, The batch_size
handling currently silently coerces invalid string specs and accepts
non-positive integers; update the block that sets self._auto_rule and
self.batch_size (around batch_size, self._auto_rule, and _compute_batch_size in
lmdb_data.py) to strictly validate inputs: accept only "auto" or "auto:<positive
int>" for string specs (parse and set self._auto_rule or raise ValueError on any
other format), and when batch_size is numeric ensure int(batch_size) > 0 (raise
ValueError otherwise), also catch and re-raise parsing errors with a clear
message so invalid configs fail fast instead of being coerced.

---

461-461: ⚠️ Potential issue | 🟠 Major

Remove unused local unique_nlocs (F841).

Line 461 assigns unique_nlocs but never uses it.

Suggested fix

-        unique_nlocs = sorted(self._nloc_groups.keys())
         nloc_info = ", ".join(
             f"{nloc}({len(idxs)})" for nloc, idxs in sorted(self._nloc_groups.items())
         )

As per coding guidelines, **/*.py: Always run ruff check . and ruff format . before committing changes or CI will fail.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` at line 461, The local variable
unique_nlocs is assigned via sorted(self._nloc_groups.keys()) but never used,
causing an F841; remove the unused assignment (delete the unique_nlocs line) or
replace its usage if it was intended to be used later in the same
function/method that contains self._nloc_groups; ensure you reference the
variable name unique_nlocs and the expression sorted(self._nloc_groups.keys())
when making the change, then run ruff check . and ruff format . before
committing.

---

848-855: ⚠️ Potential issue | 🟡 Minor

Handle empty datasets explicitly in get_test().

If the LMDB has zero frames, Line 854 calls max(...) on an empty dict and raises a cryptic error. Add an explicit empty-dataset guard.

Suggested fix

         if nloc is not None:
             if nloc not in self._nloc_groups:
                 raise ValueError(
                     f"No frames with nloc={nloc}. Available: {sorted(self._nloc_groups.keys())}"
                 )
             frame_indices = self._nloc_groups[nloc]
             natoms = nloc
+        elif not self._frames:
+            raise ValueError("LMDB dataset contains no frames")
         elif len(self._nloc_groups) == 1:
             # Uniform nloc — use all frames
             natoms = next(iter(self._nloc_groups))

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 848 - 855, get_test() can
call max(...) on an empty self._nloc_groups when the LMDB has zero frames; add
an explicit guard at the start of get_test() (or before the branching that uses
self._nloc_groups) to detect an empty dataset (e.g., if not self._frames or not
self._nloc_groups) and handle it cleanly by returning an empty result or raising
a clear exception; update the branch that sets natoms/frame_indices (the block
using self._nloc_groups, natoms, and frame_indices) to assume non-empty only
after this guard so max(...) is never called on an empty dict.

---

1030-1032: ⚠️ Potential issue | 🟠 Major

Guard frame_nlocs metadata indexing in merge path.

Line 1031 assumes src_nlocs has nframes entries. Truncated/malformed metadata raises IndexError and aborts merge.

Suggested fix

-                if src_nlocs is not None:
-                    frame_nlocs.append(int(src_nlocs[i]))
-                else:
+                if src_nlocs is not None and i < len(src_nlocs):
+                    frame_nlocs.append(int(src_nlocs[i]))
+                else:
                     frame_raw = msgpack.unpackb(raw, raw=False)
                     atype_raw = frame_raw.get("atom_types")
                     if isinstance(atype_raw, dict):

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 1030 - 1032, The code
unconditionally indexes src_nlocs[i] when building frame_nlocs in the merge
path, which can raise IndexError for truncated/malformed metadata; update the
merge logic (the block that appends to frame_nlocs using src_nlocs, referencing
variables frame_nlocs, src_nlocs, nframes, and loop index i) to guard access:
check that src_nlocs is not None and i < len(src_nlocs) before
converting/append; if missing, append a sensible default (e.g., 0) or log a
warning and continue so the merge does not abort.

deepmd/dpmodel/utils/lmdb_data.py (3)

762-765: Consider deterministic seeding for test reproducibility.

The shuffle uses np.random.default_rng() without a seed, making test results non-reproducible. While randomization is often desirable for testing, consider adding an optional seed parameter for debugging reproducibility.

♻️ Proposed enhancement

     def __init__(
         self,
         lmdb_path: str,
         type_map: list[str] | None = None,
         shuffle_test: bool = True,
+        seed: int | None = None,
         **kwargs: Any,
     ) -> None:
         ...
         # Shuffle if requested
         if shuffle_test:
-            rng = np.random.default_rng()
+            rng = np.random.default_rng(seed)
             indices = rng.permutation(len(self._frames))
             self._frames = [self._frames[i] for i in indices]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 762 - 765, The test shuffle
is non-deterministic because it calls np.random.default_rng() without a seed;
update the method (the code that checks shuffle_test and shuffles self._frames)
to accept an optional seed parameter (default None) and pass it into
np.random.default_rng(seed) so permutation(self._frames) becomes reproducible
when a seed is provided; ensure callers that invoke this behavior can pass the
seed (or leave None for non-deterministic behavior) and keep using indices =
rng.permutation(len(self._frames)) and self._frames = [self._frames[i] for i in
indices] unchanged.

---

1007-1045: Consider using context manager for source LMDB environments.

While src_env.close() is called, an exception during the loop could leave environments open. Using a context manager would ensure cleanup.

♻️ Proposed enhancement

     for src_path in src_paths:
-        src_env = _open_lmdb(src_path)
-        with src_env.begin() as txn:
-            meta = _read_metadata(txn)
-        nframes, src_fmt, natoms_per_type = _parse_metadata(meta)
-        ...
-        src_env.close()
+        with lmdb.open(src_path, readonly=True, lock=False, readahead=False, meminit=False) as src_env:
+            with src_env.begin() as txn:
+                meta = _read_metadata(txn)
+            nframes, src_fmt, natoms_per_type = _parse_metadata(meta)
+            ...

Note: Verify that lmdb.Environment supports context manager protocol (it does since lmdb 0.94).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 1007 - 1045, Wrap source LMDB
environments opened by _open_lmdb in a context manager to guarantee they are
closed on exceptions: replace the pattern "src_env = _open_lmdb(src_path); with
src_env.begin() as txn: ...; src_env.close()" with "with _open_lmdb(src_path) as
src_env: ..." and keep the existing nested "with src_env.begin() as src_txn,
dst_env.begin(write=True) as dst_txn:" block intact; remove the explicit
src_env.close() call and ensure code still references src_nlocs, src_txn,
frame_idx, and other variables unchanged.

---

245-247: Thread-safety assumption for persistent transaction.

The comment at Line 246 states the persistent transaction is "Safe because we use num_workers=0 in DataLoader." This is a critical assumption—if a caller uses num_workers > 0 with PyTorch DataLoader, the LMDB transaction will be shared across processes (via fork), which can cause corruption or hangs.

Consider adding a runtime check or explicit documentation in the class docstring warning against multi-worker usage, or lazily initializing the transaction per-worker.